import numpy as np
from numpy import random as rand
import matplotlib
from matplotlib import pyplot as plt

POP_SIZE = 10
DNA_SIZE = 10
EPOCH_NUM = 50
CROSS_RATE = 0.8
MUTATION_RATE = 0.05


def cross_and_mutation(pop):
    new_pop = []
    for father in pop:
        child = father
        if np.random.rand() <= CROSS_RATE:
            mother = pop[rand.randint(POP_SIZE)]
            cross_idx = rand.randint(low=0, high=DNA_SIZE)
            child[cross_idx:] = mother[cross_idx:]
        mutation(child)
        new_pop.append(child)
    return np.array(new_pop)

def mutation(dna):
    if rand.rand() <= MUTATION_RATE:
        mut_idx = rand.randint(0, DNA_SIZE)
        dna[mut_idx] = dna[mut_idx] ^ 1


def select(pop, fitness):
    # 使用softmax来增大适应性最强的个体被选中的概率
    fitness_copy = []
    fitness_copy.extend(fitness)
    fitness_copy -= np.max(fitness_copy)
    softmax = np.exp(fitness_copy) / np.sum(np.exp(fitness_copy))
    selection = rand.choice(np.arange(POP_SIZE), size=POP_SIZE, replace=True, p=softmax)
    return pop[selection]


def get_fitness(pop, fit):
    fitness = []
    for dna in pop:
        fitness.append(fit(dna))
    return np.array(fitness)


def print_info(pop, fitness):
    max_idx = np.argmax(fitness)
    print(pop[max_idx])
    print("fitness:{}".format(fitness[max_idx]))
    return pop[max_idx]


if __name__ == '__main__':
    pop = np.random.randint(2, size=(POP_SIZE, DNA_SIZE))

    def fit(dna):
        res = 0
        for i in range(len(dna)):
            res += dna[i] * (2 ** i)
        return res ** 2


    matplotlib.use('TKAgg')
    plt.figure()
    fitness = get_fitness(pop, fit)
    x = []
    y = []
    plt.ion()
    for i in range(EPOCH_NUM):
        pop = cross_and_mutation(pop)
        fitness = get_fitness(pop, fit)
        # print("epoch:{}, fitness avg:{}".format(i, np.average(fitness)))
        print_info(pop, get_fitness(pop, fit))
        x.append(i + 1)  # 添加 i 到 x 轴的数据中
        y.append(np.sqrt(np.max(fitness)))  # 添加 i 的平方到 y 轴的数据中
        plt.clf()  # 清除之前画的图
        plt.xlabel('epochs')
        plt.ylabel('current optimal')
        plt.plot(x, y)  # 画出当前 ax 列表和 ay 列表中的值的图形
        plt.pause(0.01)  # 暂停一秒
        plt.ioff()
        pop = select(pop, fitness)
    print_info(pop, get_fitness(pop, fit))
    plt.ioff()
    plt.plot(x, y)
    plt.show()
